from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated, Sequence
import operator
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_community.chat_models import ChatTongyi
from IPython.display import Image, display

from dotenv import load_dotenv
load_dotenv(verbose=True)

# 一个简单的langgraph示例
model = ChatTongyi(
    model="qwen-turbo",
    temperature=0,
)

# 定义state
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

# 定义Node
def _call_model(state):
    response = model.invoke(state["messages"])
    return {"messages": [response]}


if __name__ == '__main__':
    # 定义一个graph
    workflow = StateGraph(AgentState)
    # 添加一个节点
    workflow.add_node("model", _call_model)
    # 设置入口
    workflow.set_entry_point("model")
    # 设置普通边
    workflow.add_edge("model", END)
    # 编译
    app = workflow.compile()

    res = app.invoke({"messages": [HumanMessage(content="你好")]})
    print(res)
